"""

Quantum computer simulation similar to qcl. 
Usage examples:

>>> from quantsim import *
>>> heap = QHeap(5) # initializing a quantum heap
>>> print heap
(1+0j) |00000>
>>> a = heap.allocate(1) # allocating a quantum register of size 1
>>> print a
1.0 |0>

>>> Not(a) # Applying boolean Not on register a
>>> print a
1.0 |1>

>>> b = heap.allocate(2)
>>> print b
1.0 |00>
>>> Not(b[1])
>>> CNot(b[0], b[1]) # Applying conditional Not
>>> print b
1.0 |11>

>>> heap = QHeap(5)
>>> a = heap.allocate(1)
>>> Mix(a) # Applying Hadamard operator to register a
>>> print a
0.5 |0> + 0.5 |1>
>>> Mix(a) # Mix is its own inverse
>>> print a
1.0 |0>

>>> b = heap.allocate(1)
>>> Not(a)
>>> Mix(a)
>>> Not(b)
>>> Mix(b)
>>> print a
0.5 |0> + 0.5 |1>
>>> print b
0.5 |0> + 0.5 |1>
>>> print heap # a and b are completely independent
(0.5+0j) |00000> + (-0.5+0j) |00001> + (-0.5+0j) |00010> + (0.5+0j) |00011>
>>> Measure(a) # We can measure a...
1
>>> print b # ... and b stays the same
0.5 |0> + 0.5 |1>
>>> print heap
(-0.707106781187+0j) |0001> + (0.707106781187+0j) |0011>

>>> heap.reset() # Reset quantum heap to lowest eigenstate
>>> print heap
(1+0j) |00000>
>>> print a
1.0 |0>
>>> print b
1.0 |0>

>>> Not(a)
>>> Mix(a)
>>> CNot(b, a)
>>> print a
0.5 |0> + 0.5 |1>
>>> print b # a and b look the same as before
0.5 |0> + 0.5 |1>
>>> print heap # however, a and b are entangled
(0.707106781187+0j) |00000> + (-0.707106781187+0j) |00011>
>>> Measure(a) # measuring a ...
1
>>> print b # ... collapses b
1.0 |1>

>>> heap.reset()
>>> Mix(a)
>>> CPhase(3.14159, a) # Applying conditional phase shift
>>> print heap # Shifts the amplitude of |00001> by exp(-pi*j) = -1
(0.707106781187+0j) |00000> + (-0.707106781184+1.87637133737e-06j) |00001>
>>> heap.reset()
>>> Mix(a)
>>> CPhase(0.01, a) # Shifts the amplitude of |00001> by exp(0.01j)
>>> print heap
(0.707106781187+0j) |00000> + (0.707071426142+0.00707094996132j) |00001>
>>> print a
0.5 |0> + 0.5 |1>

"""

import random
from cmath import exp
from math import sqrt, log

class Qubit(object):

    def __init__(self, index, heap):
        self.index = index
        self.heap = heap
        
    def __len__(self):
        return 1

class QReg(tuple):

    def __new__(cls, heap, index, size):
        qubits = [Qubit(index + i, heap) for i in range(size)]
        return tuple.__new__(cls, qubits)

    def __init__(self, heap, index, size):
        self.heap = heap
        self.index = index

    def __str__(self):
        states = self.heap.substate(self.index, len(self))
        str_repr = self.heap.to_string(states)
        return str_repr

    def __repr__(self):
        return self.__str__()

class QHeap(object):
    
    def __init__(self, size):
        if (size <= 0):
            raise ValueError("QHeap must be initialized with positive number of Qubits")
        self.__size = size
        self.reset(size)
        self.__first_free_index = 0

    def to_string(self, states):
        str_buffer = []
        size = log(len(states), 2)
        size = int(size)
        for state, amplitude in states.iteritems():
            if abs(amplitude) < 1e-6:
                continue
            binary_repr_state = bin(state)
            list_repr_state = list(binary_repr_state[2:])
            list_repr_state.reverse()
            list_repr_state.extend("0" * (size - len(list_repr_state)))
            list_repr_state.reverse()
            str_buffer.append("%s |%s>" % (amplitude, "".join(list_repr_state)))
        return " + ".join(str_buffer)

    def __str__(self):
        return self.to_string(self.__states)

    def __repr__(self):
        return self.__str__()
    
    @property
    def states(self):
        return self.__states

    def nnot(self, index):
        if index >= self.__size:
            raise IndexError("qubit index %s out of range (max index is %s)" % (index, self.__size-1))
        if index < 0:
            raise IndexError("qubit index %s is less than zero" % index)
        bit_mask = 2 ** index
        for state in range(2 ** self.__size):
            if state & bit_mask == 0:
                self.__states[state], self.__states[state + bit_mask] = self.__states[state + bit_mask], self.__states[state]

    def cnot(self, target_index, test_indices):
        if (target_index >= self.__size):
            raise IndexError("qubit index %s is out of range (max index is %s)" % (target_index, self.__size-1))
        if (target_index < 0):
            raise IndexError("qubit index %s is less than zero" % (target_index))
        test_bit_mask = 0
        for index in test_indices:
            if index < 0 or index >= self.__size:
                raise IndexError("qubit index %s is out of range" % index)
            test_bit_mask = 2 ** index
        target_bit_mask = 2 ** target_index
        for state in range(2 ** self.__size):
            test_qubits_are_one = (state & test_bit_mask == test_bit_mask)
            target_qubit_is_zero = (state & target_bit_mask == 0)
            if test_qubits_are_one and target_qubit_is_zero:
                self.__states[state], self.__states[state + target_bit_mask] = self.__states[state + target_bit_mask], self.__states[state]

    def cphase(self, a, indices, reverse = False):
        for index in indices:
            if index >= self.__size:
                raise IndexError("qubit indices %s out of range (max indices is %s)" % (index, self.__size-1))
            if index < 0:
                raise IndexError("qubit indices %s is less than zero" % index)
        if reverse:
            a = -a
        bit_mask = 0
        for index in indices:
            bit_mask = 2 ** index
        states = self.__states
        for state, amplitude in states.iteritems():
            if state & bit_mask == bit_mask:
                states[state] = amplitude * exp(complex(0, a))

    def mix(self, index):
        if index >= self.__size:
            raise IndexError("qubit index %s out of range (max index is %s)" % (index, self.__size-1))
        if index < 0:
            raise IndexError("qubit index %s is less than zero" % index)
        states = self.__states
        new_states = {}
        bit_mask = 2 ** index
        for state, amplitude in states.iteritems():
            if state & bit_mask == 0:
                zero_state = state
                one_state = state + bit_mask
                new_states[zero_state] = new_states.get(zero_state, complex(0, 0)) + sqrt(0.5)*amplitude
                new_states[one_state] = new_states.get(one_state, complex(0, 0)) + sqrt(0.5)*amplitude
            else:
                zero_state = state - bit_mask
                one_state = state
                new_states[zero_state] = new_states.get(zero_state, complex(0, 0)) + sqrt(0.5)*amplitude
                new_states[one_state] = new_states.get(one_state, complex(0, 0)) - sqrt(0.5)*amplitude
        self.__states = new_states

    def reset(self, size = None):
        if size == None:
            size = self.__size
        num_eigenstates = 2 ** size
        eigenstates = range(num_eigenstates)
        complex_amplitudes = [complex(0,0)] * num_eigenstates
        self.__states = dict(zip(eigenstates, complex_amplitudes))
        self.__states[0] = complex(1, 0)
        

    def measure(self, index):
        # validate input
        if index >= self.__size:
            raise IndexError("qubit index %s out of range (max index is %s)" % (index, self.__size-1))
        if index < 0:
            raise IndexError("qubit index %s is less than zero" % index)

        # determine the measured state of the qubit
        target_probability = random.random()
        amplitude_sum = 0
        bit_mask = 2**index
        for state, amplitude in self.__states.iteritems():
            amplitude_sum += abs(amplitude) ** 2
            if amplitude_sum >= target_probability:
                measured_qubit_state = (state & bit_mask) >> index
                break

        # determine probability of that qubit and
        # make new dictionary only containing surviving states
        new_states = {}
        p_i = 0
        for state, amplitude in self.__states.iteritems():
            if (state & bit_mask) >> index == measured_qubit_state:
                new_states[state] = amplitude
                p_i += abs(amplitude) ** 2

        # scale new_states by 1/sqrt(p_i)
        for state, amplitude in new_states.iteritems():
            new_states[state] = amplitude / sqrt(p_i)

        # set the state to be new_states
        self.__states = new_states
        return measured_qubit_state

    def substate(self, index, size):
        if index < 0:
            raise IndexError("qubit index %s is less than zero" % index)
        if size <= 0:
            raise IndexError("substate size %s is less than zero" % size)
        if size + index > self.__size:
            raise IndexError("sum of index and size %s is out of range (max index is %s)" % (size + index, self.__size))
        substates = {}
        bit_mask = 0
        for i in range(index, index + size):
            bit_mask += 2**i
        for state, amplitude in self.__states.iteritems():
            substate = (state & bit_mask) >> index
            substates[substate] = substates.get(substate, 0) + abs(amplitude) ** 2
        p_i = 0
        for amplitude in substates.itervalues():
            p_i += abs(amplitude) ** 2
        for state, amplitude in substates.iteritems():
            substates[state] = amplitude**2 / p_i
        return substates

    def allocate(self, size):
        if size <= 0:
            raise IndexError("reg size must be greater than zero")
        if size > self.__size:
            raise MemoryError("heap out of memory")
        if self.__first_free_index + size > self.__size:
            raise MemoryError("heap out of memory")
        reg = QReg(self, self.__first_free_index, size)
        self.__first_free_index += size
        return reg
    
def CNot(target_qreg, test_qreg):
    if target_qreg.heap is not test_qreg.heap:
        raise ValueError("Both registers must be from the same heap")
    
    target_indices = []
    if len(target_qreg) == 1:
        target_indices.append(target_qreg.index)
    else:
        for target_qubit in target_qreg:
            target_indices.append(target_qubit.index)
            
    test_indices = []
    if len(test_qreg) == 1:
        test_indices.append(test_qreg.index)
    else:
        for test_qubit in test_qreg:
            test_indices.append(test_qubit.index)        
    
    heap = target_qreg.heap
    for target_index in target_indices:
        heap.cnot(target_index, test_indices)
        
def CPhase(a, reg, reverse=False):
    if len(reg) == 1:
        indices = [reg.index]
    else:
        indices = [qubit.index for qubit in reg]
    heap = reg.heap
    heap.cphase(a, indices, reverse)

def apply_method_that_takes_an_index_as_an_argument(reg, heap_method):
    if len(reg) == 1:
        heap_method(reg.index)
        return
    for qubit in reg:
        index = qubit.index
        heap_method(index)
    
def Mix(reg):
    heap = reg.heap
    heap_method = heap.mix
    apply_method_that_takes_an_index_as_an_argument(reg, heap_method)
        
def Not(reg):
    heap = reg.heap
    heap_method = heap.nnot
    apply_method_that_takes_an_index_as_an_argument(reg, heap_method)
    
def Measure(reg):
    heap = reg.heap
    heap_method = heap.measure
    apply_method_that_takes_an_index_as_an_argument(reg, heap_method)
    index = reg.index
    size = len(reg)
    reg_state = heap.substate(index, size)
    for state, amplitude in reg_state.iteritems():
        if abs(amplitude) == 1:
            return state
    assert False

if __name__ == "__main__":
    h = QHeap(2)
    a = h.allocate(2)
    Not(a)
    Mix(a[0])
    Mix(a[1])
    print a
